import os
import time
import random
import argparse

import torch 
import numpy as np
from torch.optim.lr_scheduler import MultiStepLR
from utils.earlystopping import EarlyStopping
# print(torch.cuda.is_available())

def set_seed(seed):
    torch.manual_seed(seed) # 为CPU设置随机种子
    torch.cuda.manual_seed(seed) # 为当前GPU设置随机种子
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU，为所有GPU设置随机种子
    np.random.seed(seed)  # Numpy module.
    random.seed(seed)  # Python random module.	
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def main(args):
    set_seed(args.seed)

    save_path = "checkpoints/{}/{}/".format(args.ds, args.lo)
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    if args.use_earlystopping:
        # early = EarlyStopping(patience=50, path=os.path.join(save_path, "{}_lo={}_seed={}.pt".format(args.ds, args.lo, args.seed)))
        early = EarlyStopping(patience=50)

    
    # 定义I个Trainer支持co-teaching,，默认I为2, 否则显存不够
    from trainer import Trainer
    from trainer_ensemble import TrainerEnsemble
    trainers = [ Trainer(args, id=i) for i in range(args.num_trainers)]
    ensemble_trainer = TrainerEnsemble(args, trainers)

    best_val, best_test, best_epoch = -1, -1, -1

    for epoch in range(args.ep):
        t = time.time()
        ensemble_trainer.train_each_model(epoch=epoch)

        valid_acc = ensemble_trainer.eval_all_models(mode="valid")
        test_acc  = ensemble_trainer.eval_all_models(mode="test")
        print("Using {} s".format(time.time() - t))
        print("Epoch {:>3d}, valid acc: {:.2f}, test acc: {:.2f}. ".format(epoch+1, valid_acc, test_acc))

        if args.use_earlystopping:
            if epoch >= 1:
                early(valid_acc, epoch)
            if early.early_stop:
                break

        if valid_acc > best_val:
            best_val = valid_acc
            best_epoch = epoch
            best_test = test_acc

    print("Best Epoch {:>3d}, Best valid acc: {:.2f}, test acc: {:.2f}. ".format(best_epoch, best_val, best_test))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument('--optimizer', help='optimizer\'s name', type=str, default="sgd")
    parser.add_argument('--lr', help='optimizer\'s learning rate', type=float, default=5e-2)
    parser.add_argument('--wd', help='weight decay', type=float, default=1e-3)
    parser.add_argument('--bs', help='batch size', type=int, default=256)
    parser.add_argument('--ep', help='number of epochs', type=int, default=250)
    # datasets
    parser.add_argument('--ds', help='specify a dataset', type=str, default='cifar10', required=False)
    parser.add_argument('--annotation_type', type=str, default='partial_ins', required=False)
    # training
    parser.add_argument('--mo', help='model name', type=str, default='resnet', required=False)
    parser.add_argument('--lo', help='specify a loss function', default='proden', type=str, required=False)
    parser.add_argument('--num_trainers', type=int, default=2)
    parser.add_argument('--T', type=float, default=1.0)
    # pseudo-label
    parser.add_argument('--finetune_eps', help='number of epochs', type=int, default=1)
    parser.add_argument("--use_exchange", type=int, default=1)
    parser.add_argument("--change_rate", type=float, default=0.5)
    parser.add_argument("--use_step_refine_F", type=int, default=1)
    parser.add_argument("--use_epoch_refine_F", type=int, default=0)
    parser.add_argument("--enhance_interval", type=int, default=1)
    # loss
    parser.add_argument("--loss_use_origin", type=int, default=0)
    parser.add_argument("--loss_origin_weight", type=float, default=1)
    parser.add_argument("--loss_weight_kappa", type=float, default=0.5)
    # 
    parser.add_argument('--lambda_list', nargs='+', type=float, default=[1, ])
    parser.add_argument('--beta_list', nargs='+', type=float, default=[0.7, ])
    # ema
    parser.add_argument('--use_ema', type=int, default=1)
    parser.add_argument("--ema_step", type=int, default=40000)
    parser.add_argument("--ema_alpha", type=float, default=0.997)
    parser.add_argument("--ema_mode", type=str, default="sigmoid")
    # roll window
    parser.add_argument('--use_rollwindow', type=int, default=1)
    parser.add_argument("--window_size", type=int, default=5)
    # 
    parser.add_argument('--seed', help='random seed', default=0, type=int)
    parser.add_argument('--gpu_ids', nargs='+', type=int, default=[0, ])
    parser.add_argument('--use_earlystopping', type=int, default=1)
    args = parser.parse_args()

    print(args)
    main(args)

